# ------------------------------------------------------------------------
# Copyright (c) Hitachi, Ltd. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
import argparse

import torch
from torch import nn


def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--load_path', type=str, required=True,
    )
    parser.add_argument(
        '--save_path', type=str, required=True,
    )
    parser.add_argument(
        '--dataset', type=str, default='hico',
    )
    parser.add_argument(
        '--num_queries', type=int, default=100,
    )

    args = parser.parse_args()

    return args


def main(args):
    ps = torch.load(args.load_path)

    for k in list(ps['model_state_dict'].keys()):
        ps['model']=ps['model_state_dict']
    
    del ps['model_state_dict']

    for k in list(ps['model'].keys()):
        print(k)
        if 'decoder' in k:
            ps['model'][k.replace('decoder', 'human_decoder')] = ps['model'][k].clone()


    ps['model']['query_embed.weight'] = ps['model']['query_embed.weight'].clone()[:args.num_queries]

    for i in range(3):
        ps['model']['sub_bbox_embed.layers.{}.weight'.format(i)] = ps['model']['bbox_embed.layers.{}.weight'.format(i)].clone()
        ps['model']['sub_bbox_embed.layers.{}.bias'.format(i)] = ps['model']['bbox_embed.layers.{}.bias'.format(i)].clone()
        
        ps['model']['obj_bbox_embed.layers.{}.weight'.format(i)] = ps['model']['bbox_embed.layers.{}.weight'.format(i)].clone()
        ps['model']['obj_bbox_embed.layers.{}.bias'.format(i)] = ps['model']['bbox_embed.layers.{}.bias'.format(i)].clone()
        
        del ps['model']['bbox_embed.layers.{}.weight'.format(i)]
        del ps['model']['bbox_embed.layers.{}.bias'.format(i)]

    ps['model']['obj_class_embed.weight'] = ps['model']['class_embed.weight'].clone()
    ps['model']['obj_class_embed.bias'] = ps['model']['class_embed.bias'].clone()


    del ps['model']['class_embed.weight']
    del ps['model']['class_embed.bias']




    if args.dataset == 'vcoco':
        #l = nn.Linear(ps['model']['obj_class_embed.weight'].shape[1], 1)
        #l.to(ps['model']['obj_class_embed.weight'].device)
        ps['model']['obj_class_embed.weight'] = torch.cat((ps['model']['obj_class_embed.weight'][:-1], 
                                                           #l.weight, 
                                                           ps['model']['obj_class_embed.weight'].clone()[:1],
                                                           ps['model']['obj_class_embed.weight'][-1:]))
        ps['model']['obj_class_embed.bias'] = torch.cat((ps['model']['obj_class_embed.bias'][:-1], 
                                                         #l.bias, 
                                                         ps['model']['obj_class_embed.bias'].clone()[:1],
                                                         ps['model']['obj_class_embed.bias'][-1:]))
    ############################################################################################
    print('turning.............')
    for k in list(ps['model'].keys()):
        print(k)

    torch.save(ps, args.save_path)


if __name__ == '__main__':
    args = get_args()
    main(args)
